#!/usr/bin/python3.8
# -*- coding: utf-8 -*-
#
# Copyright (C) 2022 Liu Yuting All Rights Reserved #
# @Time    : 2022/10/25 15:44
# @Author  : Liu Yuting
# @Email   : 1972456540@qq.com
# @File    : Foil_Parametric_fitting.py
# @Software: PyCharm

from scipy.optimize import leastsq
from scipy.special import factorial
import numpy as np
from matplotlib import pyplot as plt
import re




def MakeFunCST(n, trailbias, surface, b=1, control=1):
    """ CST翼型参数方程外函数

    Args:
        n (int): Bernstein 多项式次数
        trailbias(list): 翼型后缘偏移，单位:m
            [x, y], x是上表面后缘偏移，y是下表面后缘偏移
        surface (int): 方程构造对象，主要为区别trailbias
            If surface = 1, 构造翼型上表面的方程
            If surface = 0, 构造翼型下表面的方程
        b (float): 翼型弦长，单位:m
        control (int): 拟合翼型控制参数
            If control = 1 (default), 对应NACA系列的圆头尖尾翼型
            If control = 2 , 对应圆头圆尾的椭圆翼型
            If control = 3 , 对应尖头尖尾的翼型

    return:
        FunCST (function): CST参数方程函数FunCST(alist, x)

    Require:
        import numpy as np
        from scipy.special import factorial

    """
    # 控制系数
    if control == 1:
        N1, N2 = 0.5, 1
    elif control == 2:
        N1, N2 = 0.5, 0.5
    else:
        N1, N2 = 1, 1

    def FunCST(alist, x):
        """ CST翼型参数方程构造

        Args:
            alist (list): 控制参数 [a1, a2, a3, ...]，alist长度为 n+1
            x (list): 表达式未知数x，表示翼型上横坐标的值

        Returns:
            y (ndarry): CST参数方程

        """

        if len(alist) != (n + 1):
            ex = Exception("the length of alist is not equal to n+1")
            raise ex
        # 控制参数赋值
        for i in range(1, n + 2):
            exec("a{0} = {1}".format(i, alist[i - 1]))

        # 类别函数
        x = np.array(x, dtype=float)
        C = (x / b) ** N1 * (1 - x / b) ** N2

        # 形状函数
        i = np.arange(0, n + 1, dtype=float)
        K = factorial(n) / (factorial(i) * factorial(n - i))
        a = np.array(alist, dtype=float)
        S = [sum(a * K * ((xi / b) ** i * (1 - xi / b) ** (n - i))) for xi in x]

        # 纵坐标表达式
        if surface == 1:  # 上表面
            y = C * S + x / b * trailbias[0] / b
        else:  # 下表面
            y = C * S + x / b * trailbias[1] / b

        return y

    return FunCST


def MakeError(MakeFun, method, para, surface):
    """ 残差函数外函数

    Args:
        MakeFun (function): 翼型参数化方程外函数
            MakeFunCST: CST参数化方法
        method (str): 拟合方法
            "CST" : 采用CST方法拟合
        para (Dict): 拟合方法相应参数
            If method is "CST": para有4对键值，分别为：
                "n" (int): Bernstein多项式次数
                "b" (float): 翼型弦长，单位:m
                "trailbias" (list): 翼型后缘偏移，单位:m，一般上下表面值不同
                    [x, y], x是上表面后缘偏移，y是下表面后缘偏移
                "control" (int): 拟合翼型控制参数，通常为1
                    If control = 1 (default), 对应NACA系列的圆头尖尾翼型
                    If control = 2 , 对应圆头圆尾的椭圆翼型
                    If control = 3 , 对应尖头尖尾的翼型
        surface (int): 方程构造对象，主要为区别trailbias
            If surface = 1, 构造翼型上表面的方程
            If surface = 0, 构造翼型下表面的方程

    return:
        Error (function): 残差函数Error(alist, x, y)

    """

    if method == "CST":
        makefun = MakeFun(para["n"], para["trailbias"], surface, para["b"], para["control"])

    def Error(alist, x, y):
        """ 残差函数

        Args:
            alist (list): 控制参数
                If method is "CST": alist = [a1, a2, a3, ...]，alist长度为 n+1
            x (list): 表达式未知数x，表示翼型横坐标的值

        Returns:
            残差 (ndarry)

        """

        return makefun(alist, x) - y

    return Error


def FoilFit(FoilData, method, para, makeerror, makefun, printresult=True):
    """ 翼型参数化拟合

    Args:
        FoilData (from FoilLoader): 翼型数据
        method (str): 拟合方法
            "CST" : 采用CST方法拟合
        para (Dict): 拟合方法相应参数，具体参数请参考 MakeError 函数
        makeerror (function): 残差外函数
        makefun (funciton): 拟合参数方程外函数
        printresult (bool): 打印参数化列表和RSS

    Return:
        controlpara (Dict): 翼型拟合控制参数
            ControlParaUp (list): 翼型上表面参数列表
            ControlParaLow (list): 翼型下表面参数列表

    Require:
        import numpy as np
        from matplotlib import pyplot as plt
        from scipy.optimize import leastsq

    """

    if method == "CST":
        n = para["n"]
        # 参数初始值，与alist维度相同
        PInitUp = np.full(n + 1, 0.5)
        PInitLow = np.full(n + 1, -0.5)

    # 拟合
    RssFun = makeerror(makefun, method="CST", para=para, surface=1)
    ControlParaUp = leastsq(RssFun, PInitUp, args=(FoilData["UpX"], FoilData["UpY"]))[0]
    RssFun = makeerror(makefun, method="CST", para=para, surface=0)
    ControlParaLow = leastsq(RssFun, PInitLow, args=(FoilData["LowX"], FoilData["LowY"]))[0]

    # 打印结果
    if printresult:
        # 残差平方和Rss
        RssUp = sum(np.array(RssFun(ControlParaUp, FoilData["UpX"], FoilData["UpY"])) ** 2)
        RssLow = sum(np.array(RssFun(ControlParaLow, FoilData["LowX"], FoilData["LowY"])) ** 2)

        print("上表面拟合参数列表：\n{}".format(ControlParaUp))
        print("上表面拟合残差平方和RSS：{}".format(RssUp))
        print("下表面拟合参数列表：\n{}".format(ControlParaLow))
        print("下表面拟合残差平方和RSS：{}".format(RssLow))

    controlpara = {"ControlParaUp": list(ControlParaUp), "ControlParaLow": list(ControlParaLow)}
    return controlpara


def FitPaint(controlpara, makefun, method, para, oridata, n):
    """ 绘制原翼型和拟合后翼型的轮廓

    Args:
        controlpara (Dict): 翼型拟合控制参数
            ControlParaUp (list): 翼型上表面参数列表
            ControlParaLow (list): 翼型下表面参数列表
        makefun (funciton): 拟合参数方程外函数
        method (str): 拟合方法
            "CST" : 采用CST方法拟合
        para (Dict): 拟合方法相应参数，具体参数请参考 MakeError 函数
        oridata (Dict): 原翼型数据
            If oridata is None: UpX和Low取自分布函数

    Require:
        from matplotlib import pyplot as plt
        import numpy as np
    """

    # 设置绘图风格
    # plt.style.use('ggplot')
    # 中文和负号的正常显示
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False

    OnlyPaintFit = False
    if oridata is None:
        OnlyPaintFit = True
        Num = 200
        # 横坐标采样点分布函数
        z1 = [-2.26220338e-17, 1.57541715e-14, -6.26292265e-13, -1.24799869e-09, -1.92308714e-09, 6.23473358e-05, -3.35345218e-07, 8.87161535e-07]
        f1 = np.poly1d(z1)
        UpX = LowX = [0] + list(f1(np.linspace(0, 200, Num-1).tolist()))
        oridata = {"Name": "XXX",
                   "Num": Num,
                   "UpX": UpX,
                   "LowX": LowX
                   }

    if method == "CST":
        funcstup = makefun(para["n"], para["trailbias"], 1, para["b"], para["control"])
        funcstlow = makefun(para["n"], para["trailbias"], 0, para["b"], para["control"])
        title = "CST Method Bernstein={}".format(para["n"])

    fitdataUpX = oridata["UpX"]
    fitdataUpY = funcstup(controlpara["ControlParaUp"], fitdataUpX)
    fitdataLowX = oridata["LowX"]
    fitdataLowY = funcstlow(controlpara["ControlParaLow"], fitdataLowX)

    color0 = ["lightskyblue", "lightgreen", "thistle"]
    color = ["dodgerblue", "green", "fuchsia"]

    if OnlyPaintFit is True:
        plt.plot(fitdataUpX, fitdataUpY, c="blue", label='CSTfit foil', linewidth=1)
        plt.plot(fitdataLowX, fitdataLowY, c="blue", linewidth=1)
    else:
        plt.plot(oridata["UpX"], oridata["UpY"], c=color0[n], label=oridata["Name"] + "(Original)", linewidth=2.5)
        plt.plot(oridata["LowX"], oridata["LowY"], c=color0[n], linewidth=2.5)
        plt.plot(fitdataUpX, fitdataUpY, "--", dashes=(5, 5), c=color[n], label=oridata["Name"] + "(CST)", linewidth=2.5)
        plt.plot(fitdataLowX, fitdataLowY, "--", dashes=(5, 5), c=color[n], linewidth=2.5)

    plt.axis("equal")
    plt.rcParams.update({'font.size': 25})  # 设置图例字体大小
    plt.legend(loc='upper right')  # 定义图标所处位置，这里表示右上
    plt.title(title, fontdict={'weight': 'bold', 'size': 35})
    plt.xticks(fontsize=25)
    plt.yticks(fontsize=25)
    # plt.savefig("foil.svg", format='svg')
    # plt.show()


if __name__ == "__main__":
    # CST拟合例子
    path = r".\C72.dat"
    from Foil_Data_Process import FoilLoader
    FoilData = FoilLoader(path, normalize=True)
    para = {"n": 6, "b": FoilData["Chord"], "trailbias": [FoilData["UpY"][-1], FoilData["LowY"][-1]], "control": 1}
    controlpara = FoilFit(FoilData, method="CST", para=para, makeerror=MakeError, makefun=MakeFunCST)
    # 绘图
    FitPaint(controlpara, MakeFunCST, method="CST", para=para, oridata=FoilData, n=2)
    plt.show()



